# -*- coding: utf-8 -*-
"""
plot_face_arm.py  — 固定解析 + 收敛分析 + 恢复4张原图
规则：
- 数据从第6行(1-based)开始为数值
- 横坐标 = Time = 第2列(1-based) → 索引 1
- X/Y/Z  = 第3~5列(1-based) → 索引 2,3,4
输出：
- analysis_result.txt  （--save 可改名/改路径）
- error_with_convergence.png （误差范数曲线）
- face_xyz.png, arm_xyz.png, face_arm_3d.png, face_arm_overlay.png （恢复的4张图）
"""

import argparse
import os
import re
from typing import Tuple

import numpy as np
import matplotlib.pyplot as plt


# ---------- split utils ----------
def _split_keep_empty(line: str):
    """保留空单元格分割：优先 TAB，其次逗号/分号，否则连续空白。"""
    s = line.rstrip("\r\n")
    if "\t" in s:
        return s.split("\t")
    if "," in s:
        return s.split(",")
    if ";" in s:
        return s.split(";")
    return re.split(r"\s+", s.strip())


# ---------- readers ----------
def read_ly_time_xyz(path: str) -> Tuple[np.ndarray, np.ndarray, np.ndarray]:
    """
    读取 .ly 数据（硬规则）：
    - 数据从第6行开始
    - Frame# = 第1列（若非数值则用1..N）
    - t = 第2列（Time），xyz = 第4~6列（注意你当前脚本就是用的 row[3],row[4],row[5]）
    返回 t:(N,), xyz:(N,3), frames:(N,)
    """
    with open(path, "r", encoding="utf-8", errors="ignore") as f:
        lines = f.readlines()

    if len(lines) < 6:
        raise RuntimeError(f"文件行数不足 6 行: {path}")

    data_rows = []
    for ln in lines[5:]:  # 第6行起
        ln = ln.strip()
        if not ln:
            continue
        toks = _split_keep_empty(ln)
        data_rows.append(toks)

    def to_float(x, default=np.nan):
        try:
            return float(str(x).strip())
        except Exception:
            return default

    def to_int_or(x, default):
        try:
            return int(str(x).strip())
        except Exception:
            return default

    t_list, x_list, y_list, z_list, frames = [], [], [], [], []
    auto_frame = 1
    for row in data_rows:
        if len(row) <= 5:  # 至少 6 列
            continue

        # 帧号（第1列）；若非数值，用自增帧号兜底
        fid = to_int_or(row[0], auto_frame)
        if fid == auto_frame:
            auto_frame += 1
        else:
            auto_frame = fid + 1

        t = to_float(row[1])   # Time（第2列）
        x = to_float(row[3])   # X（第4列）
        y = to_float(row[4])   # Y（第5列）
        z = to_float(row[5])   # Z（第6列）

        if np.isfinite(t) and np.isfinite(x) and np.isfinite(y) and np.isfinite(z):
            frames.append(fid)
            t_list.append(t); x_list.append(x); y_list.append(y); z_list.append(z)

    t = np.asarray(t_list, dtype=float)
    xyz = np.vstack([x_list, y_list, z_list]).T.astype(float)
    frames = np.asarray(frames, dtype=int)
    return t, xyz, frames


# ---------- align by nearest time ----------
def align_by_time(t1: np.ndarray, v1: np.ndarray,
                  t2: np.ndarray, v2: np.ndarray) -> Tuple[np.ndarray, np.ndarray, np.ndarray]:
    """最近邻时间配对对齐两路数据，返回 (t_aligned, v1_aligned, v2_aligned)。"""
    if len(t1) == 0 or len(t2) == 0:
        raise ValueError("时间序列为空")

    ref_is_1 = len(t1) <= len(t2)  # 选短的做参考
    if ref_is_1:
        t_ref, v_ref = t1, v1
        t_find, v_find = t2, v2
    else:
        t_ref, v_ref = t2, v2
        t_find, v_find = t1, v1

    idx_sorted = np.argsort(t_find)
    t_find_s = t_find[idx_sorted]
    v_find_s = v_find[idx_sorted]

    idx_closest = []
    for tr in t_ref:
        pos = np.searchsorted(t_find_s, tr)
        cand = []
        if pos > 0:
            cand.append(pos - 1)
        if pos < len(t_find_s):
            cand.append(pos)
        if not cand:
            continue
        best = min(cand, key=lambda k: abs(t_find_s[k] - tr))
        idx_closest.append(best)

    t_aligned = []
    v_ref_a, v_find_a = [], []
    for i_ref, i_find in zip(range(len(t_ref)), idx_closest):
        t_aligned.append(0.5 * (t_ref[i_ref] + t_find_s[i_find]))
        v_ref_a.append(v_ref[i_ref])
        v_find_a.append(v_find_s[i_find])

    t_aligned = np.asarray(t_aligned, dtype=float)
    v_ref_a = np.asarray(v_ref_a, dtype=float)
    v_find_a = np.asarray(v_find_a, dtype=float)

    if ref_is_1:
        return t_aligned, v_ref_a, v_find_a
    else:
        return t_aligned, v_find_a, v_ref_a

def detect_motion_start_by_err_bend(
    t: np.ndarray,
    err_norm: np.ndarray,
    neg_slope_thresh: float = 0.0,   # <=0 则自适应（推荐）
    flat_slope_thresh: float = 0.0,  # <=0 则自适应（推荐）
    pre_window: float = 0.25,        # 需要在此时间长度内基本平台
    hold_sec: float = 0.12,          # 下降斜率需要持续的时间
    smooth_N: int = 7,               # 平滑窗口
    baseline_sec: float = 0.40,      # 用于估计噪声的“起始基线时长”
    flat_scale: float = 1.8,         # 平台阈值 = flat_scale * 噪声
    drop_scale: float = 3.0          # 下降阈值 = drop_scale * 噪声
) -> int:
    """
    在误差范数曲线上找“拐点”作为运动起点：
    1) 平滑 err_norm -> e
    2) 计算导数 dE/dt -> deriv（右对齐到 t[i]）
    3) 用起始 baseline_sec 内的导数作为噪声基线，得到自适应阈值
    4) 找首个 j 满足：
       - j 之前 pre_window 内 |deriv| 的 *均值* <= flat_slope_thresh
       - j 到 j+hold_sec 内 deriv 的 *均值* <= -neg_slope_thresh
    返回：起点索引（对齐 t 的索引）。找不到则返回 -1。
    """
    n = len(t)
    if n < 5 or n != len(err_norm):
        return -1

    # 1) 平滑
    e = err_norm.astype(float)
    if smooth_N and smooth_N > 1:
        k = int(max(1, smooth_N))
        pad = np.pad(e, (k // 2, k - 1 - k // 2), mode="edge")
        e = np.convolve(pad, np.ones(k) / k, mode="valid")  # 长度 n

    # 2) 一阶导（导数 deriv[i] 近似对应 t[i]）
    dt = np.diff(t)
    if not np.any(dt > 0):
        return -1
    safe_dt = dt.copy()
    safe_dt[safe_dt <= 1e-6] = np.median(safe_dt[safe_dt > 0])
    de = np.diff(e)
    deriv_raw = de / safe_dt               # 长度 n-1
    deriv = np.concatenate([[deriv_raw[0]], deriv_raw])  # 右对齐为长度 n，与 t 对齐

    med_dt = float(np.median(np.diff(t)))
    need_hold = max(1, int(round(hold_sec / med_dt)))
    need_pre  = max(1, int(round(pre_window / med_dt)))
    need_base = max(need_pre, int(round(baseline_sec / med_dt)))

    # 3) 基线噪声（只看开头 need_base 段的导数）
    base = deriv[:min(len(deriv), need_base)]
    if len(base) < 3:
        return -1
    # 用稳健统计（中位数绝对偏差）估计噪声量级
    med = float(np.median(base))
    mad = float(np.median(np.abs(base - med))) + 1e-12
    noise = 1.4826 * mad  # ≈ σ

    # 自适应阈值
    if flat_slope_thresh <= 0:
        flat_slope_thresh = max(1e-6, flat_scale * noise)
    if neg_slope_thresh <= 0:
        neg_slope_thresh = max(1e-6, drop_scale * noise)

    # 4) 搜索 j：前段平台、后段持续下降（均值判定更鲁棒）
    # 注意索引范围：确保两段窗口都完整
    for j in range(need_pre, len(deriv) - need_hold):
        pre_mean = float(np.mean(np.abs(deriv[j-need_pre:j])))
        if pre_mean > flat_slope_thresh:
            continue
        post_mean = float(np.mean(deriv[j:j+need_hold]))
        if post_mean <= -abs(neg_slope_thresh):
            return j  # 对齐到 t[j]

    return -1

# ---------- motion start detection (NEW) ----------
def detect_motion_start_by_speed(t: np.ndarray,
                                 xyz: np.ndarray,
                                 speed_thresh: float = 1.0,
                                 hold_sec: float = 0.15,
                                 smooth_N: int = 5) -> int:
    """
    基于末端(arm)速度幅值检测运动起点：
    - 速度 = ||Δxyz|| / Δt
    - 平滑：窗口 smooth_N 的滑动平均
    - 连续 hold_sec 时间超过 speed_thresh 判为“开始运动”
    返回：开始的索引（对应 t 的索引）。若未超过阈值，返回 0。
    """
    n = len(t)
    if n < 3 or n != len(xyz):
        return 0

    dt = np.diff(t)
    if not np.any(dt > 0):
        return 0
    dt[dt <= 1e-6] = np.median(dt[dt > 0])

    v_inst = np.linalg.norm(np.diff(xyz, axis=0), axis=1) / dt  # 单位同 XYZ/秒

    # 平滑
    if smooth_N and smooth_N > 1:
        k = int(max(1, smooth_N))
        pad = np.pad(v_inst, (k - 1, 0), mode="edge")
        v_smooth = np.convolve(pad, np.ones(k) / k, mode="valid")
    else:
        v_smooth = v_inst

    # 连续阈值时长 -> 需要的点数
    need = max(1, int(round(hold_sec / np.median(np.diff(t)))))
    for i in range(0, len(v_smooth) - need + 1):
        if np.all(v_smooth[i:i + need] >= speed_thresh):
            return i  # v_smooth 相对 t 少 1 个点，直接用 i 对应 t[i]
    return 0


# ---------- convergence ----------
def find_convergence_time(t: np.ndarray,
                          err: np.ndarray,
                          stable_sec: float = 0.5,
                          rel_tol: float = 0.05,
                          abs_tol: float = 1.0,
                          start_idx: int = 0):
    """
    误差范数收敛判定：
    - 末尾~0.8s中位数为 e_final
    - 带宽 = max(abs_tol, rel_tol*|e_final|)
    - 从 start_idx 开始找首个点，使后续 stable_sec 内误差均在 [e_final ± 带宽]
    返回 (t_conv_from_start, idx_conv, e_final)；未收敛：t_conv=NaN, idx_conv=-1
    """
    if len(t) != len(err) or len(t) < 3:
        return np.nan, -1, float(err[-1]) if len(err) else np.nan

    dt = np.median(np.diff(t))
    dt = float(dt) if dt > 0 else 1.0 / 90.0
    win_final = max(30, int(round(0.8 / dt)))
    e_final = float(np.median(err[-win_final:])) if len(err) >= win_final else float(np.median(err))

    band = max(abs_tol, rel_tol * abs(e_final))
    need = max(10, int(round(stable_sec / dt)))

    s = max(0, int(start_idx))
    for i in range(s, len(err) - need + 1):
        seg = err[i:i + need]
        if np.all(np.abs(seg - e_final) <= band):
            return float(t[i] - t[s]), i, e_final

    return np.nan, -1, e_final


# ---------- main ----------
def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("face_file", type=str, help=".ly: face")
    ap.add_argument("arm_file",  type=str, help=".ly: arm")
    ap.add_argument("--save", type=str, default="analysis_result.txt", help="结果文本输出路径")
    ap.add_argument("--stable-sec", type=float, default=0.5, help="稳定判定需要持续的时间(s)")
    ap.add_argument("--rel-tol", type=float, default=0.05, help="相对带宽（相对最终误差）")
    ap.add_argument("--abs-tol", type=float, default=1.0, help="绝对带宽（单位同XYZ）")
    ap.add_argument("--no-show", action="store_true", help="不弹出图窗，仅保存")
    ap.add_argument("--frame-start", type=int, default=None, help="按 Frame# 选择：起始帧（含）")
    ap.add_argument("--frame-end",   type=int, default=None, help="按 Frame# 选择：结束帧（含）")

    # ===== NEW: 起点检测参数 =====
    ap.add_argument("--start-mode", choices=["speed", "bend", "none"], default="bend",
                    help="收敛时间起点：speed=arm速度检测，bend=误差曲线拐点，none=从首帧")
    ap.add_argument("--speed-thresh", type=float, default=1.0,
                    help="速度阈值（单位=XYZ单位/秒）")
    ap.add_argument("--start-hold", type=float, default=0.15,
                    help="速度超过阈值需持续的时间(s)")
    # bend（误差拐点）参数
    ap.add_argument("--bend-neg-slope", type=float, default=8.0,
                    help="bend：下降斜率阈值（单位=误差单位/秒；<=0 启用自适应）")
    ap.add_argument("--bend-hold", type=float, default=0.12,
                    help="bend：负斜率需持续的时间(s)")
    ap.add_argument("--bend-flat-slope", type=float, default=2.0,
                    help="bend：平台斜率阈值")
    ap.add_argument("--bend-pre-window", type=float, default=0.44,
                    help="bend：判定前的平台期时长(s)")

    args = ap.parse_args()

    out_dir = os.path.dirname(args.save) or "."
    os.makedirs(out_dir, exist_ok=True)

    # 读取
    t_face, face_xyz, f_frames = read_ly_time_xyz(args.face_file)
    t_arm,  arm_xyz,  a_frames = read_ly_time_xyz(args.arm_file)

    # 可选：按帧号选择（含端点）
    if (args.frame_start is not None) or (args.frame_end is not None):
        def select_by_frame(t, xyz, frames, f0, f1):
            if f0 is None:
                f0 = int(frames.min())
            if f1 is None:
                f1 = int(frames.max())
            mask = (frames >= f0) & (frames <= f1)
            if not np.any(mask):
                raise RuntimeError(f"选择的帧区间 [{f0}, {f1}] 无数据")
            return t[mask], xyz[mask], frames[mask]

        t_face, face_xyz, f_frames = select_by_frame(
            t_face, face_xyz, f_frames, args.frame_start, args.frame_end)
        t_arm,  arm_xyz,  a_frames = select_by_frame(
            t_arm, arm_xyz, a_frames, args.frame_start, args.frame_end)
        print(f"[face] frames: {f_frames[0]}..{f_frames[-1]} (kept {len(f_frames)})")
        print(f"[arm ] frames: {a_frames[0]}..{a_frames[-1]} (kept {len(a_frames)})")

    # 画 4 张图（与原版保持一致）
    # 图1：face
    plt.figure()
    plt.plot(t_face, face_xyz[:, 0], label="X_face")
    plt.plot(t_face, face_xyz[:, 1], label="Y_face")
    plt.plot(t_face, face_xyz[:, 2], label="Z_face")
    plt.xlabel("Time (s)"); plt.ylabel("value")
    plt.title("face: X/Y/Z"); plt.legend(); plt.tight_layout()
    plt.savefig(os.path.join(out_dir, "face_xyz.png")); plt.close()

    # 图2：arm
    plt.figure()
    plt.plot(t_arm, arm_xyz[:, 0], label="X_arm")
    plt.plot(t_arm, arm_xyz[:, 1], label="Y_arm")
    plt.plot(t_arm, arm_xyz[:, 2], label="Z_arm")
    plt.xlabel("Time (s)"); plt.ylabel("value")
    plt.title("arm: X/Y/Z"); plt.legend(); plt.tight_layout()
    plt.savefig(os.path.join(out_dir, "arm_xyz.png")); plt.close()

    # 图3：3D 叠加
    try:
        from mpl_toolkits.mplot3d import Axes3D  # noqa: F401
        fig = plt.figure(); ax = fig.add_subplot(111, projection="3d")
        ax.plot(face_xyz[:, 0], face_xyz[:, 1], face_xyz[:, 2], label="face")
        ax.plot(arm_xyz[:, 0],  arm_xyz[:, 1],  arm_xyz[:, 2],  label="arm")
        ax.set_xlabel("X"); ax.set_ylabel("Y"); ax.set_zlabel("Z")
        ax.set_title("3D Trajectories"); ax.legend(); plt.tight_layout()
        ax.set_xlim(240, 280)
        plt.savefig(os.path.join(out_dir, "face_arm_3d.png"))
        # plt.close()
        plt.show() 
    except Exception as e:
        print("3D 绘图失败（可忽略）：", e)

    # 图4：叠加曲线
    plt.figure()
    plt.plot(t_face, face_xyz[:, 0], label="face.X")
    plt.plot(t_arm,  arm_xyz[:, 0],  label="arm.X",  linestyle="--")
    plt.plot(t_face, face_xyz[:, 1], label="face.Y")
    plt.plot(t_arm,  arm_xyz[:, 1],  label="arm.Y",  linestyle="--")
    plt.plot(t_face, face_xyz[:, 2], label="face.Z")
    plt.plot(t_arm,  arm_xyz[:, 2],  label="arm.Z",  linestyle="--")
    plt.xlabel("Time (s)"); plt.ylabel("value")
    plt.title("face vs arm — X/Y/Z (overlay)")
    plt.legend(ncol=2); plt.tight_layout()
    plt.savefig(os.path.join(out_dir, "face_arm_overlay.png")); plt.close()

    # 对齐后做误差/收敛
    t, face_a, arm_a = align_by_time(t_face, face_xyz, t_arm, arm_xyz)
    err_xyz = arm_a - face_a
    err_norm = np.linalg.norm(err_xyz, axis=1)

    # ===== 起点检测 =====
    if args.start_mode == "speed":
        start_idx = detect_motion_start_by_speed(
            t, arm_a,
            speed_thresh=args.speed_thresh,
            hold_sec=args.start_hold
        )
    elif args.start_mode == "bend":
        start_idx = detect_motion_start_by_err_bend(
            t, err_norm,
            neg_slope_thresh=args.bend_neg_slope,
            flat_slope_thresh=args.bend_flat_slope,
            pre_window=args.bend_pre_window,
            hold_sec=args.bend_hold
        )
        # 找不到拐点则回退到 speed 法
        if start_idx in (-1, 0):
            fallback = detect_motion_start_by_speed(
                t, arm_a,
                speed_thresh=args.speed_thresh,
                hold_sec=args.start_hold
            )
            if fallback > 0:
                print(f"[hint] bend 未找到拐点，已回退 speed 起点: idx={fallback}")
                start_idx = fallback
            else:
                print("[hint] bend 未找到拐点且 speed 起点也未满足，使用首帧(0)")
                start_idx = 0
    else:  # none
        start_idx = 0


    # 收敛判定（从 start_idx 起计时）
    t_conv, idx_conv, e_final = find_convergence_time(
        t, err_norm,
        stable_sec=args.stable_sec,
        rel_tol=args.rel_tol,
        abs_tol=args.abs_tol,
        start_idx=start_idx
    )

    # 最终误差分量（末尾0.8s中位数）
    dt = np.median(np.diff(t)) if len(t) > 1 else 1.0/90.0
    win_final = max(30, int(round(0.8 / dt)))
    tail_slice = slice(max(0, len(err_xyz) - win_final), len(err_xyz))
    e_final_xyz = np.median(err_xyz[tail_slice, :], axis=0)
    e_final_norm = float(np.linalg.norm(e_final_xyz))

    # 文本结果
    lines = []
    lines.append(f"Convergence judge: stable_sec={args.stable_sec}, rel_tol={args.rel_tol}, abs_tol={args.abs_tol}")
    lines.append(f"start_mode={args.start_mode}, start_idx={start_idx}")
    if args.start_mode == "speed":
        lines.append(f"speed_thresh={args.speed_thresh}, start_hold={args.start_hold}")
    if idx_conv < 0 or not np.isfinite(t_conv):
        lines.append("Converged: NO")
        lines.append(f"final_error_norm   : {e_final_norm:.6f}")
        lines.append(f"final_error_xyz    : {e_final_xyz[0]:.6f}, {e_final_xyz[1]:.6f}, {e_final_xyz[2]:.6f}")
        lines.append(f"frame-start        : {args.frame_start:.2f}")
        lines.append(f"frame-end          : {args.frame_end:.2f}")
    else:
        lines.append("Converged: YES")
        lines.append(f"convergence_time(s): {t_conv:.6f}  (from motion start)")
        lines.append(f"index_at_converged : {idx_conv}")
        lines.append(f"final_error_norm   : {e_final_norm:.6f}")
        lines.append(f"final_error_xyz    : {e_final_xyz[0]:.6f}, {e_final_xyz[1]:.6f}, {e_final_xyz[2]:.6f}")
        lines.append(f"frame-start        : {args.frame_start:.2f}")
        lines.append(f"frame-end          : {args.frame_end:.2f}")



    with open(args.save, "w", encoding="utf-8") as f:
        f.write("\n".join(lines))
    print("\n".join(lines))
    print(f"\n分析结果已写入: {args.save}")

    # 误差范数曲线（横轴=Time）
    err_img = os.path.join(out_dir, "error_with_convergence.png")
    plt.figure(figsize=(10, 6))
    plt.plot(t, err_norm, label="||arm - face||")

    # NEW: 起点与收敛两条辅助线
    t0_plot = t[start_idx] if len(t) > 0 else 0.0
    plt.axvline(x=t0_plot, linestyle=":", label="Motion start")
    if idx_conv >= 0:
        plt.axvline(x=t[idx_conv], linestyle="--",
                    label=f"Converged @ {(t[idx_conv]-t0_plot):.2f}s")

    plt.xlabel("Time (s)"); plt.ylabel("Error Norm (same unit as XYZ)")
    plt.title("Face vs Arm Error over Time")
    plt.legend(); plt.grid(True, alpha=0.3); plt.tight_layout()
    plt.savefig(err_img)
    if not args.no_show:
        plt.show()
    plt.close()
    print(f"误差曲线已保存: {err_img}")


if __name__ == "__main__":
    main()
